package com.aloudmeta.grpc.support;

import io.grpc.ForwardingServerCallListener;
import io.grpc.Metadata;
import io.grpc.ServerCall;
import io.grpc.ServerCallHandler;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.rpc.protocol.grpc.interceptors.ServerInterceptor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.SocketAddress;

import static io.grpc.Grpc.TRANSPORT_ATTR_REMOTE_ADDR;
import static org.apache.dubbo.common.constants.CommonConstants.PROVIDER;

/**
 * huaixin 2022/2/21 3:25 PM
 */
@Activate(group = {PROVIDER}, order=5400)
public class HeaderServerInterceptor implements ServerInterceptor {

    private static final Logger LOG = LoggerFactory.getLogger(HeaderServerInterceptor.class);

    private static Metadata.Key<String> GRPC_TENANT_ID =
            Metadata.Key.of("TENANT_ID", Metadata.ASCII_STRING_MARSHALLER);

    @Override
    public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
            ServerCall<ReqT, RespT> call, Metadata headers, ServerCallHandler<ReqT, RespT> next) {
        String tenantId = headers.get(GRPC_TENANT_ID);
        LOG.info("tenantId ====== {}", tenantId);
        call.setCompression("zstd");
        LOG.info("header received from client:" + headers);
        SocketAddress socketAddress = call.getAttributes().get(TRANSPORT_ATTR_REMOTE_ADDR);
        LOG.info("socketAddress ====== {}", socketAddress);
        return new HeaderServerInterceptor.HeaderServerCallListener<>(next.startCall(call, headers), call);
    }

    private class HeaderServerCallListener<ReqT, RespT> extends ForwardingServerCallListener.SimpleForwardingServerCallListener<ReqT> {

        private final ServerCall<ReqT, RespT> call;

        protected HeaderServerCallListener(final ServerCall.Listener<ReqT> delegate,
                                                        final ServerCall<ReqT, RespT> call) {
            super(delegate);
            this.call = call;
        }

        @Override
        // Unary calls error out here
        public void onHalfClose() {
            super.onHalfClose();
        }

    }
}
